# -*- coding: utf-8 -*-
# @Time : 2022/6/14 上午11:28
# @Author : wyf
# @FileName: 异步获取视频.py
# @Email: yunfei.wang@innvote.com
# path: /home/ywt/Downloads/video
import os
import re
import time

import aiofiles
import requests
import asyncio, aiohttp

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from aiohttp import ClientTimeout, ClientError
from base import logger


def get_m3u8(url, file_name):
    req = requests.get(url)
    with open(f"{file_name}.m3u8", "w", encoding="utf-8") as v:
        v.write(req.text)
    return req.text


async def read_m3u8(file_name="video.m3u8"):
    with open(file_name, mode="r", encoding="utf-8") as f:
        tasks = []
        for index, line in enumerate(f):
            if line.startswith("#"):
                continue
            ts_name = f"{index}.ts"
            ts_url = line.strip()
            tasks.append(asyncio.create_task(download_ts(ts_url, ts_name)))
        await wait_task(tasks)


async def wait_task(tasks):
    """ 在这里进行添加协程任务, 并获取任务的结果, 如果结果为失败那么就重新添加到任务中 """
    await asyncio.wait(tasks)
    err_tasks = []
    for task in tasks:
        if task.result() is not None:
            _url, _name = task.result()
            err_tasks.append(asyncio.create_task(download_ts(_url, _name)))
    if not err_tasks:
        print("获取视频结束~~~")
        return
    print("下载错误的TS")
    await wait_task(err_tasks)


async def download_ts(ts_url, ts_name):
    timeout = ClientTimeout(total=30)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        try:
            async with session.get(ts_url) as res:
                # print(f'[INFO]请求成功{ts_url, ts_name}')
                if res.status != 200:
                    logger.warning(f'请求错误{res.status}--{ts_url, ts_name}')
                    return ts_url, ts_name
                if os.path.exists(video_path) is False:
                    os.mkdir(video_path)
                async with aiofiles.open(f"{video_path}/{ts_name}", mode="wb") as v:
                    await v.write(await res.content.read())
        except asyncio.TimeoutError as e:
            logger.warning(f'请求超时错误{ts_url, ts_name}')
            return ts_url, ts_name


async def des_ts():
    """
    from Crypto.Util.Padding import pad
    encrypted_data = pad(encrypted_data, 16)
    :return:
    """
    try:
        cipher = AES.new(b"04fff35febefbf9c", AES.MODE_CBC, IV=b'0000000000000000')
        for p in os.listdir(video_path):
            async with aiofiles.open(os.path.join(video_path, p), 'rb') as file:
                # TODO './news/' 为解密后的视频存放位置
                async with aiofiles.open('./news/' + p, "wb") as fp:
                    content = await file.read()
                    # 获取长度
                    encrypted_data_len = len(content)
                    # 判断当前的数据长度是不是16的倍数
                    if encrypted_data_len % 16 != 0:
                        # 把长度不是16的倍数的显示出来
                        # 变为16的倍数
                        content = pad(content, 16)
                    plain_data = cipher.decrypt(content)
                    if os.path.exists('news') is False:
                        os.mkdir('news')
                    await fp.write(plain_data)

    except Exception as e:
        print(p)
        print(e)


def merge_file():
    ts_files = os.listdir(video_path)
    ts_files.sort(key=lambda x: int(x.split(".")[0]))
    shell_str = "+".join(ts_files)
    video_name = "video.mp4"
    # 删除ts视频
    strs = 'copy /b ' + shell_str + f' {video_name}' + '\n' + 'del *.ts'
    # 不删除ts视频
    # strs = 'copy /b ' + shell_str + f' {names}'
    f = open("combined.cmd", 'w')
    f.write(strs)
    f.close()


def run():
    # 下载 TS 文件
    get_m3u8(m3u8_url, m3u8_file_name)
    # asyncio.get_event_loop().run_until_complete(read_m3u8())
    asyncio.run(read_m3u8())
    # 视频合并
    merge_file()

    # # 视频解密
    # asyncio.run(des_ts())


if __name__ == '__main__':
    m3u8_file_name = "video.m3u8"
    video_path = "/home/ywt/Downloads/video"
    m3u8_url = "https://new.iskcd.com/20210830/uX0yMdri/1300kb/hls/index.m3u8"

    s_time = time.time()
    run()
    print(time.time() - s_time)

